package com.light.dp.flink.sink;

import com.alibaba.fastjson.JSONObject;
import com.light.dp.flink.config.dto.ColumnMapping;
import com.light.dp.flink.config.dto.JDBCDatabaseConfig;
import com.light.dp.flink.config.dto.TableMapping;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.springframework.util.CollectionUtils;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

/**
 * 存量备份处理器
 */
public class StockSinkHandler extends RichSinkFunction<String> {
    private static final int MAX_RETRIES = 3;
    private Connection connection;
    private List<PreparedStatement> preparedStatements;
    private JDBCDatabaseConfig targetConfig;
    private List<TableMapping> tableMappings;

    public StockSinkHandler(JDBCDatabaseConfig targetConfig, List<TableMapping> tableMappings) {
        this.targetConfig = targetConfig;
        this.tableMappings = tableMappings;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        String url = String.format("jdbc:%s://%s:%d/%s", targetConfig.getJdbcDriver().getJdbcDrive(), targetConfig.getHostname(), targetConfig.getPort(), targetConfig.getDatabase());
        connection = DriverManager.getConnection(url, targetConfig.getUsername(), targetConfig.getPassword());
        preparedStatements = new ArrayList<>();

        for (TableMapping tableMapping : tableMappings) {
            List<ColumnMapping> columnMappings = tableMapping.getColumnMappings();
            if (CollectionUtils.isEmpty(columnMappings)) {
                // 如果没有手动设置，就使用全表字段映射
                String targetTable = tableMapping.getTargetTable();
                List<String> columns = getTableColumns(targetTable);
                List<String> primaryKeys = getTablePrimaryKeys(targetTable);

                columnMappings = new ArrayList<>();
                for (String column : columns) {
                    boolean isPrimaryKey = primaryKeys.contains(column);
                    columnMappings.add(new ColumnMapping(column, column, isPrimaryKey));
                }
                tableMapping.setColumnMappings(columnMappings);
            }

            StringBuilder columns = new StringBuilder();
            StringBuilder values = new StringBuilder();
            StringBuilder updateClause = new StringBuilder();
            List<String> primaryKeys = new ArrayList<>();
            for (int i = 0; i < columnMappings.size(); i++) {
                ColumnMapping mapping = columnMappings.get(i);
                String targetColumn = mapping.getTargetColumn();
                if (i > 0) {
                    columns.append(", ");
                    values.append(", ");
                }
                columns.append(targetColumn);
                values.append("?");
                if (mapping.getIsTargetPrimaryKey()) {
                    primaryKeys.add(targetColumn);
                } else {
                    if (updateClause.length() > 0) {
                        updateClause.append(", ");
                    }
                    updateClause.append(targetColumn).append(" = VALUES(").append(targetColumn).append(")");
                }
            }
            String insertQuery = "INSERT INTO " + tableMapping.getTargetTable() + " (" + columns + ") VALUES (" + values + ")";
            if (updateClause.length() > 0) {
                insertQuery += " ON DUPLICATE KEY UPDATE " + updateClause;
            }
            preparedStatements.add(connection.prepareStatement(insertQuery));
        }
    }

    @Override
    public void invoke(String value, SinkFunction.Context context) throws Exception {
//        System.out.println(value);
        JSONObject json = JSONObject.parseObject(value);
        for (int i = 0; i < tableMappings.size(); i++) {
            TableMapping tableMapping = tableMappings.get(i);
            List<ColumnMapping> columnMappings = tableMapping.getColumnMappings();
            PreparedStatement preparedStatement = preparedStatements.get(i);
            int retries = 0;
            while (retries < MAX_RETRIES) {
                try {
                    for (int j = 0; j < columnMappings.size(); j++) {
                        ColumnMapping mapping = columnMappings.get(j);
                        String sourceColumn = mapping.getSourceColumn();
                        Object val = json.getJSONObject("after").get(sourceColumn);
                        if ("id".equals(sourceColumn)) {
                            System.out.println("mysql：" + val.toString());
                        }
                        // 检查是否为 create_time 列，如果是则进行时间戳转换
                        if ("create_time".equals(sourceColumn)) {
                            if (val != null) {
                                long timestamp = Long.parseLong(val.toString());
                                // 将毫秒时间戳转换为 java.sql.Timestamp
                                Timestamp sqlTimestamp = new Timestamp(timestamp);
                                val = sqlTimestamp;
                            }
                        }

                        preparedStatement.setObject(j + 1, val);
                    }
                    preparedStatement.executeUpdate();
                    break;
                } catch (SQLException e) {
                    // 任务失败重试三次
                    retries++;
                    if (retries == MAX_RETRIES) {
                        throw e;
                    }
                    Thread.sleep(1000);
                }
            }
        }
    }

    @Override
    public void close() throws Exception {
        for (PreparedStatement preparedStatement : preparedStatements) {
            if (preparedStatement != null) {
                preparedStatement.close();
            }
        }
        if (connection != null) {
            connection.close();
        }
    }


    private List<String> getTableColumns(String tableName) throws SQLException {
        List<String> columns = new ArrayList<>();
        Statement statement = connection.createStatement();
        ResultSet resultSet = statement.executeQuery("SELECT * FROM " + tableName + " LIMIT 0");
        int columnCount = resultSet.getMetaData().getColumnCount();
        for (int i = 1; i <= columnCount; i++) {
            columns.add(resultSet.getMetaData().getColumnName(i));
        }
        resultSet.close();
        statement.close();
        return columns;
    }


    private List<String> getTablePrimaryKeys(String tableName) throws SQLException {
        List<String> primaryKeys = new ArrayList<>();
        String query = "SELECT COLUMN_NAME " +
                "FROM information_schema.KEY_COLUMN_USAGE " +
                "WHERE TABLE_SCHEMA = ? " +
                "AND TABLE_NAME = ? " +
                "AND CONSTRAINT_NAME = 'PRIMARY'";
        PreparedStatement preparedStatement = connection.prepareStatement(query);
        preparedStatement.setString(1, targetConfig.getDatabase());
        preparedStatement.setString(2, tableName);
        ResultSet resultSet = preparedStatement.executeQuery();
        while (resultSet.next()) {
            primaryKeys.add(resultSet.getString("COLUMN_NAME"));
        }
        resultSet.close();
        preparedStatement.close();
        return primaryKeys;
    }
}
